package com.hyperf.core.shardingjdbc;

import org.apache.shardingsphere.api.sharding.standard.PreciseShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.core.rule.ShardingRule;
import org.apache.shardingsphere.core.rule.TableRule;
import org.apache.shardingsphere.underlying.common.rule.DataNode;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * 数据库分库策略
 * <p>
 * 参数结构
 * availableTargetNames ["ds_1", "ds_2"]
 * shardingValue {logicTableName:"t_order", columnName:"order_id", value:0}
 * </p>
 *
 * @author lizhaoyang
 */
@SuppressWarnings("all")
public class PreciseShardingDatabaseAlgorithm implements PreciseShardingAlgorithm {

    private final static String DATABASE_NAME_UNDERLINE = "_";

    private Map<String, Integer> logicTableNameActualDataNodesMap = new ConcurrentHashMap<>();

    private ShardingRule shardingRule;

    @Override
    public String doSharding(Collection availableTargetNames, PreciseShardingValue shardingValue) {
        int databaseIndex = getDatabaseIndex(availableTargetNames.size(), shardingValue.getLogicTableName(), shardingValue.getValue());
        for (Object availableTargetName : availableTargetNames) {
            String databaseName = String.valueOf(availableTargetName);
            if (databaseName.endsWith(DATABASE_NAME_UNDERLINE + databaseIndex)) {
                return databaseName;
            }
        }

        throw new UnsupportedOperationException();
    }

    public void setShardingRule(ShardingRule shardingRule) {
        this.shardingRule = shardingRule;
    }

    /**
     * 计算sharding对应的库索引
     *
     * @param databaseSize   分库数量
     * @param logicTableName 逻辑表名
     * @param columnValue    分片键value
     * @return 库索引
     */
    private int getDatabaseIndex(int databaseSize, String logicTableName, Object columnValue) {
        int tableSize = getActualDataNodeSize(logicTableName);
        int sumSize = databaseSize * tableSize;
        return (Math.abs(columnValue.hashCode()) % sumSize) / tableSize;
    }

    /**
     * 获取分表数量
     *
     * @param logicTableName
     * @return
     */
    private int getActualDataNodeSize(String logicTableName) {
        if (logicTableNameActualDataNodesMap.containsKey(logicTableName)) {
            return logicTableNameActualDataNodesMap.get(logicTableName);
        } else {
            if (shardingRule == null) {
                throw new UnsupportedOperationException();
            }

            int actualDataNodeSize = shardingRule
                    .getTableRule(logicTableName)
                    .getActualDataNodes()
                    .stream()
                    .map(DataNode::getTableName)
                    .collect(Collectors.toSet())
                    .size();

            logicTableNameActualDataNodesMap.put(logicTableName, actualDataNodeSize);

            return actualDataNodeSize;
        }
    }
}
